from PyQt5 import QtWidgets, QtGui, QtCore
from ui2 import Ui_MainWindow
from PyQt5.QtWidgets import *
import datetime
import time
from PyQt5.QtGui import QFont
import cv2
import numpy as np
from PIL import Image
from yolo import YOLO
from waste_category import main as type_inf
from PyQt5 import QtCore
QtCore.QCoreApplication.setAttribute(QtCore.Qt.AA_EnableHighDpiScaling)

class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
    def __init__(self, parent=None):
        super(MainWindow, self).__init__(parent)
        self.setupUi(self)


        self.label_type.setText('')
        self.label_pro.setText('')
        self.textEdit.setText('')

        # 设置字体大小
        font = QFont()
        font.setPointSize(12)
        self.textEdit.setFont(font)


        #日期更新

        today = str(datetime.datetime.now().date())

        self.label_12.setText(today.split('-')[0] + '/' + today.split('-')[1])
        self.label_13.setText(today.split('-')[2])


        self.yolo=YOLO()


        self.timer_camera = QtCore.QTimer()  # 定义定时器，用于控制显示视频的帧率 循环倒计时
        self.timer_video = QtCore.QTimer()
        self.cap = cv2.VideoCapture()  # 视频流
        self.CAM_NUM = 0  # 为0时表示视频流来自笔记本内置摄像头

        self.timer_video.timeout.connect(self.display_video_frame)  # 若定时器结束，则调用show_camera()
        self.timer_camera.timeout.connect(self.show_camera)  # 若定时器结束，则调用show_camera()

        self.pushButton_pic.clicked.connect(self.select_img)  # 选择图像检测
        self.pushButton_video.clicked.connect(self.select_video)  # 选择视频检测
        self.pushButton_camera.clicked.connect(self.button_open_camera_clicked)  # 选择摄像头检测


        self.name_dict = {
            'xie': '废旧鞋子', 'jidanke': '鸡蛋壳','shuibei': '水杯', 'xia': '虾壳',
            'qingcai': '青菜', 'canjinzhi': '餐巾纸','dao': '刀', 'wanou': '毛绒玩偶',
            'xiguapi': '西瓜皮', 'chongdianbao': '充电宝', 'suliaodai': '塑料袋', 'qiaokeli': '巧克力',
            'wenduji': '温度计', 'naicha': '奶茶', 'zhiniaoku': '纸尿裤', 'diliao': '火锅底料',
            'guo': '铁锅', 'baozhi': '报纸杂志','jidankedddaa': '鸡蛋壳','mifan': '米饭',
            'kouxiangtang': '口香糖','zhibei': '纸杯','dengpao': '灯泡','yantou': '烟头',
            'yao': '药片','chazuo': '插座','baoxianmo': '保鲜膜','bao': '包','jiuping': '酒瓶',
            'neizang': '内脏','yilaguan': '易拉罐','wazi': '袜子','huoji': '打火机','dingzi': '铁钉',
            'xiangjiaopi': '香蕉皮','mianbao': '面包','yaqian': '牙签','kouzhao': '口罩'
        }


    def select_video(self):
        self.timer_video.stop()  # 关闭定时器
        self.timer_camera.stop()  # 关闭定时器
        self.cap.release()  # 释放视频流
        self.label_show.clear()  # 清空视频显示区域
        self.label_show.setStyleSheet("background-color: rgb(0, 0, 0);\n"
                                      "border-radius: 20px;")
        self.pushButton_video.setText('选择视频')


        self.file_path, _ = QFileDialog.getOpenFileName(None, '选择视频文件', '', 'Video Files (*.mp4 *.avi)')

        if self.file_path:
            self.pushButton_video.setText('关闭/选择视频')
            self.camera = cv2.VideoCapture(self.file_path)
            self.timer_video.start(30)
        else:
            self.label_type.setText('')
            self.label_pro.setText('')
            self.textEdit.setText('')
            self.label_4.setStyleSheet("background-color: rgb(229, 229, 229);")

    def display_video_frame(self):

        flag, self.image = self.camera.read()  # 从视频流中读取

        if flag==True:
            show = cv2.resize(self.image, (400, 310))  # 把读到的帧的大小重新设置为 640x480
            cv2.imwrite('./jietu.jpg', show)
            image = Image.open('./jietu.jpg')
            r_image = self.yolo.detect_image(image, crop=False, count=False)
            r_image.save('./res2.png')
            self.label_show.setStyleSheet("image: url(./res2.png)")  # 将检测出的图片放到界面框中
            self.show_data()

        else:
            self.timer_video.stop()  # 关闭定时器

    def select_img(self):
        self.timer_video.stop()  # 关闭定时器
        self.timer_camera.stop()  # 关闭定时器
        self.cap.release()  # 释放视频流
        self.label_show.clear()  # 清空视频显示区域

        self.img_path, _ = QFileDialog.getOpenFileName(None, 'open img', '', "*.png;*.jpg;;All Files(*)")

        if self.img_path:
            print(self.img_path)
            image = Image.open(self.img_path)
            r_image = self.yolo.detect_image(image, crop=False, count=False)
            r_image=r_image.resize((400,310))
            r_image.save('res3.png')
            self.label_show.setStyleSheet("image: url(./res3.png)")  #将检测出的图片放到界面框中

            self.show_data()

    def button_open_camera_clicked(self):
        self.timer_video.stop()  # 关闭定时器
        self.label_show.clear()  # 清空视频显示区域
        self.label_show.setStyleSheet("background-color: rgb(0, 0, 0);\n"
                                      "border-radius: 20px;")
        if self.timer_camera.isActive() == False:  # 若定时器未启动

                flag = self.cap.open(self.CAM_NUM)  # 参数是0，表示打开笔记本的内置摄像头，参数是视频文件路径则打开视频
                if flag == False:  # flag表示open()成不成功
                    msg = QtWidgets.QMessageBox.warning(self, 'warning', "请检查相机于电脑是否连接正确",
                                                                    buttons=QtWidgets.QMessageBox.Ok)
                else:
                    self.pushButton_camera.setText("关闭摄像头")
                    self.timer_camera.start(30)  # 定时器开始计时30ms，结果是每过30ms从摄像头中取一帧显示

        else:
            self.timer_camera.stop()  # 关闭定时器
            self.cap.release()  # 释放视频流
            self.label_show.clear()  # 清空视频显示区域
            self.label_show.setStyleSheet("background-color: rgb(0, 0, 0);\n"
                                          "border-radius: 20px;")


            self.pushButton_camera.setText('摄像头检测')
            self.label_type.setText('')
            self.label_pro.setText('')
            self.textEdit.setText('')
            self.label_4.setStyleSheet("background-color: rgb(229, 229, 229);")




    def show_camera(self):

        flag, self.image = self.cap.read()  # 从视频流中读取
        show = cv2.resize(self.image, (400, 310))  # 把读到的帧的大小重新设置为 640x480
        cv2.imwrite('./jietu.jpg', show)
        image = Image.open('./jietu.jpg')
        r_image = self.yolo.detect_image(image, crop=False, count=False)
        r_image.save('./res2.png')
        self.label_show.setStyleSheet("image: url(./res2.png)")  #将检测出的图片放到界面框中
        self.show_data()

    def show_data(self):

        with open('res.txt','r')as fb:
            data=fb.readlines()

        if data:
            for num,i in enumerate(data):
                self.label_type.setText(self.name_dict[i.strip().split('+')[0]])
                self.label_pro.setText(i.strip().split('+')[1]+'%')
                self.textEdit.setText(type_inf(i.strip().split('+')[0])[2])
                if type_inf(i.strip().split('+')[0])[1]=='其他垃圾':
                    self.label_4.setStyleSheet("image: url(ui/qita.png)")
                elif type_inf(i.strip().split('+')[0])[1]=='有害垃圾':
                    self.label_4.setStyleSheet("image: url(ui/youhai.png)")
                elif type_inf(i.strip().split('+')[0])[1]=='厨余垃圾':
                    self.label_4.setStyleSheet("image: url(ui/chuyu.png)")
                elif type_inf(i.strip().split('+')[0])[1]=='可回收物':
                    self.label_4.setStyleSheet("image: url(ui/kehuishou.png)")


        else:
            self.label_type.setText('')
            self.label_pro.setText('')
            self.textEdit.setText('')
            self.label_4.setStyleSheet("background-color: rgb(229, 229, 229);")





if __name__ == "__main__":
    import sys

    app = QtWidgets.QApplication(sys.argv)
    mainWindow = MainWindow()
    mainWindow.show()
    sys.exit(app.exec_())